Tracking down the upstream story

Metal already has the sister bug — fixed in NNlib v0.9.32

The same failure is reported on Metal as
FluxML/Flux.jl#2633 and
JuliaGPU/Metal.jl#673.
@CarloLucibello reduced it on Metal to:

using NNlib, Metal
W, x, b = Mtl.(rand(Float32, 3, 2)), Mtl.(rand(Float32, 2, 2)), Mtl.(rand(Float32, 3))

NNlib._fast_broadcast!(tanh ∘ (+),               W*x, b)  # ✅ OK
NNlib._fast_broadcast!(NNlib.tanh_fast ∘ (+),    W*x, b)  # ❌ InvalidIRError

@christiangnrd narrowed it further: removing sign(x) from a hand-rolled
copy of tanh_fast makes it compile. The fix that landed in NNlib v0.9.32
(NNlib#666) is a
single-line NNlibMetalExt:

@device_override NNlib.tanh_fast(x) = Base.FastMath.tanh_fast(x)

CUDA has the same bug — but no upstream issue or extension fix

I searched FluxML/NNlib.jl, JuliaGPU/CUDA.jl, JuliaGPU/GPUCompiler.jl,
FluxML/Flux.jl, and LuxDL/Lux.jl. Nobody has filed it for CUDA. There's
no NNlibCUDAExt tanh_fast override either — the existing CUDA ext
only carries softplus etc. activation comments and conv_bias_act!,
not the polynomial-tanh wrapper the Metal ext has.

So we are most likely the first to trip over this on CUDA. The reason
it didn't surface earlier on main is probably that the only
Dense(_, _, tanh) GPU paths in this repo are the four DeepSplitting
tests — the rest of the suite uses relu or identity.

Reduced CUDA MWE (would need GPU hardware)

using CUDA, NNlib

a = CUDA.rand(Float32, 5)
b = CUDA.rand(Float32, 5)

# ✅ OK — Base.tanh, GPU compiler uses the device intrinsic
broadcast(tanh ∘ (+), a, b)

# ❌ InvalidIRError on `var"#_#103"(kw::Base.Pairs, c::ComposedFunction, x...)`
broadcast(NNlib.tanh_fast ∘ (+), a, b)

Equivalently, what the failing tests actually trigger:

using CUDA, Flux
m = Dense(2 => 3, tanh) |> Flux.gpu
x = CUDA.rand(Float32, 2, 4)
m(x)         # forward fails with InvalidIRError

Why tanh_fast ∘ + fails but tanh ∘ + works

NNlib.tanh_fast(x::Float32) ends in
ifelse(x2 < 66f0, x*(n/d), sign(x)). The sign branch makes the body
non-trivially type-stable inside the Broadcasted{ComposedFunction{...}}
closure, and the GPU compiler then escapes to the dynamic kwsorter
(::ComposedFunction)(x...; kw...) at operators.jl:1041 instead of
inlining the call. With Base.tanh the body inlines straight to the
device tanh intrinsic, the kwsorter dispatch never has to be resolved
dynamically, and the kernel compiles.

Where the fix really belongs

Upstream NNlib: add an NNlibCUDAExt @device_override mirroring the
Metal one. Single line:

@device_override NNlib.tanh_fast(x) = Base.FastMath.tanh_fast(x)

I'll file an issue + PR against FluxML/NNlib.jl once the workaround
here is confirmed green on the GPU runner. Until that lands and ships,
the per-array-type fast_act opt-out in this PR is the smallest
in-package fix; it's exactly the shape NNlib designed fast_act's 2nd
positional arg for (activations.jl#L897-L906).

🤖 Generated with Claude Code

GPU run on 1c0769b confirmed the fast_act opt-out resolved the four InvalidIRErrors — the GPU summary went from 222 passed, 4 errored to 227 passed, 1 failed. The new failure is just Aqua's piracy check correctly flagging the override:

Possible type-piracy detected:
[1] fast_act(::typeof(tanh), ::CUDACore.CuArray) @ HighDimPDE src/HighDimPDE.jl:29

That's intentional — we're adding a method to NNlib's fast_act for ::CuArray, which is exactly the per-array-type opt-out shape NNlib designed fast_act's 2nd positional arg for. Pushed 712fea1 whitelisting it via Aqua.test_piracies(HighDimPDE, treat_as_own = [NNlib.fast_act]). Once that goes green I'll file the upstream NNlib issue + PR proposing the equivalent NNlibCUDAExt @device_override.